package com.azh.server.vm;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import com.azh.server.datam.logger.Logger;

/**
 * @author zhaozhenhang <zhaozhenhang@kuaishou.com>
 * Created on 2022-04-30
 * 维护了一个依赖等待图，以进行死锁检测
 */
public class LockTable {

    // 某个txId已经获得资源的uid列表
    private Map<Long, List<Long>> tx2u;

    // uid被某个txId持有
    private Map<Long, Long> u2Tx;

    // 正在等待uid的txId列表
    private Map<Long, List<Long>> wait;

    // 正在等待资源的txId的锁
    private Map<Long, Lock> waitLock;

    // 正在等待的uid
    private Map<Long, Long> waitU;

    private Lock lock;

    public LockTable() {
        tx2u = new HashMap<>();
        u2Tx = new HashMap<>();
        wait = new HashMap<>();
        waitLock = new HashMap<>();
        waitU = new HashMap<>();
        lock = new ReentrantLock();
    }

    public Lock add(long txId, long uid) throws Exception {
        lock.lock();
        try {
            if (isInList(tx2u, txId, uid)) {
                return null;
            }
            if (!u2Tx.containsKey(uid)) {
                u2Tx.put(uid, txId);
                putIntoList(tx2u, txId, uid);
                return null;
            }
            waitU.put(txId, uid);
            putIntoList(wait, txId, uid);
            if (hasDeadLock()) {
                waitU.remove(txId);
                removeFromList(wait, uid, txId);
                throw new RuntimeException("Deadlock!");
            }
            Lock l = new ReentrantLock();
            l.lock();
            waitLock.put(txId, l);
            return l;
        }finally {
            lock.unlock();
        }
    }

    public void remove(long txId) {
        lock.lock();
        try {
            List<Long> l = tx2u.get(txId);
            if(l != null) {
                while(l.size() > 0) {
                    Long uid = l.remove(0);
                    selectNewXID(uid);
                }
            }
            waitU.remove(txId);
            tx2u.remove(txId);
            waitLock.remove(txId);

        } finally {
            lock.unlock();
        }
    }

    // 从等待队列中选择一个xid来占用uid
    private void selectNewXID(long uid) {
        u2Tx.remove(uid);
        List<Long> l = wait.get(uid);
        if(l == null) return;
        assert l.size() > 0;

        while(l.size() > 0) {
            long xid = l.remove(0);
            if(!waitLock.containsKey(xid)) {
                continue;
            }
            u2Tx.put(uid, xid);
            Lock lo = waitLock.remove(xid);
            waitU.remove(xid);
            lo.unlock();
            break;
        }

        if (l.size() == 0) {
            wait.remove(uid);
        }
    }

    private Map<Long, Integer> txIdStamp;
    private int stamp;

    private void removeFromList(Map<Long, List<Long>> listMap, long uid0, long uid1) {
        List<Long> l = listMap.get(uid0);
        if(l == null) return;
        Iterator<Long> i = l.iterator();
        while(i.hasNext()) {
            long e = i.next();
            if(e == uid1) {
                i.remove();
                break;
            }
        }
        if(l.size() == 0) {
            listMap.remove(uid0);
        }
    }

    private boolean hasDeadLock() {
        txIdStamp = new HashMap<>();
        stamp = 1;
        for (long txId : tx2u.keySet()) {
            Integer s = txIdStamp.get(txId);
            if (s != null && s > 0) {
                continue;
            }
            stamp++;
            if (dfs(txId)) {
                return true;
            }
        }
        return false;
    }

    private boolean dfs(long txId) {
        Integer stp = txIdStamp.get(txId);
        if(stp != null && stp == stamp) {
            return true;
        }
        if(stp != null && stp < stamp) {
            return false;
        }
        txIdStamp.put(txId, stamp);

        Long uid = waitU.get(txId);
        if(uid == null) return false;
        Long x = u2Tx.get(uid);
        assert x != null;
        return dfs(x);
    }

    private boolean isInList(Map<Long, List<Long>> listMap, long uid0, long uid1) {
        List<Long> l = listMap.get(uid0);
        if (l == null) {
            return false;
        }
        Iterator<Long> i = l.iterator();
        while (i.hasNext()) {
            long e = i.next();
            if (e == uid1) {
                return true;
            }
        }
        return false;
    }

    private void putIntoList(Map<Long, List<Long>> listMap, long uid0, long uid1) {
        if (!listMap.containsKey(uid0)) {
            listMap.put(uid0, new ArrayList<>());
        }
        listMap.get(uid0).add(0, uid1);
    }


}
